#!/usr/bin/python
"""
Author: Jimmy Saw
Last update: 08-06-2012
Description: This script runs BLASTn of sequences then retrieves top n hits as instructed.
Usage example: dissertation_BlastnRetrieveTopHits.py test.fasta 10 blastp 1e-5
"""

import sys
import re
from Bio import SeqIO
from Bio import Entrez
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML

Entrez.email = 'jimmysaw@gmail.com'

def blast(sequences, blastprog, database, maxevalue):
#    maxevalue = 0.0001
    for s in sequences:
        result_handle = NCBIWWW.qblast(blastprog, database, s.seq, 
          expect=maxevalue, filter=None)
        save_file = open(s.id + ".xml", "w")
        save_file.write(result_handle.read())
        save_file.close()
        result_handle.close()
        print "Done with BLAST for: ", s.id

def parseblastn(seqlist):
    accessions = []
    for seq in seqlist:
        xml_file = seq + ".xml"
        xf = open(xml_file, "rU")
        r = NCBIXML.parse(xf)
        results_rec = r.next()
        for hit in results_rec.alignments[0:maxhits]:
#            print seq, hit.accession, hit.hit_def
            accessions.append(hit.accession)
    return accessions

def retrieveseqs(acclist):
    sequences = []
    outfile = "top_" + str(maxhits) + "_neighbors.fasta"
    for acc in acclist:
        handle = Entrez.efetch(db="nucleotide", id=acc, rettype="fasta", 
          retmode="text")
        tmpseq = SeqIO.read(handle, "fasta")
        sequences.append(tmpseq)
    SeqIO.write(sequences, outfile, "fasta")

if __name__ == "__main__":
    seqslist = []
    seqids = []
    seqs = SeqIO.parse(sys.argv[1], "fasta")
    maxhits = int(sys.argv[2])
    blasttype = sys.argv[3]
    maxeval = sys.argv[4]
    for seq in seqs:
        seqslist.append(seq)
        seqids.append(seq.id)
    if blasttype == "blastn":
        blast(seqslist, "blastn", "nt", maxeval)
    elif blasttype == "blastp":
        blast(seqslist, "blastp", "nr", maxeval)
    elif blasttype == "blastx":
        blast(seqslist, "blastx", "nr", maxeval)
    else:
        print "Choose either: blastn, blastp, or blastx"
    accs = parseblastn(seqids)
    retrieveseqs(accs)

